import torch
import math


def positional_encoding(seq_len, d_model):
    # 生成空矩阵
    PE = torch.zeros(seq_len, d_model)
    # 生成位置索引 [0, 1, 2, ...]
    position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
    # 计算不同频率的分母项
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

    # 偶数维：sin
    PE[:, 0::2] = torch.sin(position * div_term)
    # 奇数维：cos
    PE[:, 1::2] = torch.cos(position * div_term)

    return PE


if __name__ == '__main__':
    pe = positional_encoding(10, 8)
    print(pe.shape)  # (10, 8)
    print(pe)
